package datastructure.redblacktree;

/**
 * @Author ZhangCuirong
 * @Date 2025/7/31 17:05
 * @description:
 */
import java.util.ArrayList;
import java.util.List;

public class RedBlackTree<T extends Comparable<T>> {
    // 节点颜色定义
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    // 节点内部类
    public class Node {
        public T key;           // 键值
        public Node left;       // 左子节点
        public Node right;      // 右子节点
        Node parent;     // 父节点
        public boolean color;   // 颜色

        public Node(T key, boolean color) {
            this.key = key;
            this.color = color;
            this.left = nil;
            this.right = nil;
            this.parent = nil;
        }
    }

    // 哨兵节点，所有叶子节点和空节点都指向它
    public final Node nil;
    // 根节点
    public Node root;

    // 构造函数
    public RedBlackTree() {
        // 初始化哨兵节点为黑色
        nil = new Node(null, BLACK);
        nil.left = nil;
        nil.right = nil;
        nil.parent = nil;
        root = nil;
    }

    /**
     * 左旋操作
     */
    private void leftRotate(Node x) {
        Node y = x.right;  // y是x的右子节点
        x.right = y.left;  // 将y的左子树设为x的右子树

        if (y.left != nil) {
            y.left.parent = x;  // 更新y左子树的父节点
        }

        y.parent = x.parent;  // 更新y的父节点

        // 根据x是否为根节点处理不同情况
        if (x.parent == nil) {
            root = y;
        } else if (x == x.parent.left) {
            x.parent.left = y;
        } else {
            x.parent.right = y;
        }

        y.left = x;  // 将x设为y的左子节点
        x.parent = y;  // 更新x的父节点
    }

    /**
     * 右旋操作
     */
    private void rightRotate(Node y) {
        Node x = y.left;  // x是y的左子节点
        y.left = x.right;  // 将x的右子树设为y的左子树

        if (x.right != nil) {
            x.right.parent = y;  // 更新x右子树的父节点
        }

        x.parent = y.parent;  // 更新x的父节点

        // 根据y是否为根节点处理不同情况
        if (y.parent == nil) {
            root = x;
        } else if (y == y.parent.right) {
            y.parent.right = x;
        } else {
            y.parent.left = x;
        }

        x.right = y;  // 将y设为x的右子节点
        y.parent = x;  // 更新y的父节点
    }

    /**
     * 插入节点
     */
    public void insert(T key) {
        Node z = new Node(key, RED);  // 新节点默认为红色
        Node y = nil;
        Node x = root;

        // 找到插入位置
        while (x != nil) {
            y = x;
            if (z.key.compareTo(x.key) < 0) {
                x = x.left;
            } else {
                x = x.right;
            }
        }

        z.parent = y;  // 设置z的父节点

        // 处理插入位置
        if (y == nil) {
            root = z;  // 树为空，z成为根节点
        } else if (z.key.compareTo(y.key) < 0) {
            y.left = z;
        } else {
            y.right = z;
        }

        z.left = nil;  // 左右子节点设为哨兵
        z.right = nil;
        z.color = RED;  // 新节点为红色

        // 修复红黑树性质
        insertFixup(z);
    }

    /**
     * 插入后修复红黑树性质
     */
    private void insertFixup(Node z) {
        // 当父节点为红色时需要修复（父节点为黑色则无需修复）
        while (z.parent.color == RED) {//父结点为红色
            if (z.parent == z.parent.parent.left) {//父结点是祖父结点的左结点
                Node y = z.parent.parent.right;  // 叔节点

                // 情况1：叔节点为红色
                if (y.color == RED) {
                    z.parent.color = BLACK;
                    y.color = BLACK;
                    z.parent.parent.color = RED;
                    z = z.parent.parent;  // 向上移动继续检查
                } else {
                    // 情况2：叔节点为黑色，且z是右子节点
                    if (z == z.parent.right) {
                        z = z.parent;
                        leftRotate(z);
                    }
                    // 情况3：叔节点为黑色，且z是左子节点
                    z.parent.color = BLACK;
                    z.parent.parent.color = RED;
                    rightRotate(z.parent.parent);
                }
            } else {
                // 镜像情况：父节点是祖父节点的右子节点
                Node y = z.parent.parent.left;  // 叔节点

                // 情况1：叔节点为红色
                if (y.color == RED) {
                    z.parent.color = BLACK;
                    y.color = BLACK;
                    z.parent.parent.color = RED;
                    z = z.parent.parent;  // 向上移动继续检查
                } else {
                    // 情况2：叔节点为黑色，且z是左子节点
                    if (z == z.parent.left) {
                        z = z.parent;
                        rightRotate(z);
                    }
                    // 情况3：叔节点为黑色，且z是右子节点
                    z.parent.color = BLACK;
                    z.parent.parent.color = RED;
                    leftRotate(z.parent.parent);
                }
            }
        }
        // 确保根节点始终为黑色
        root.color = BLACK;
    }

    /**
     * 查找节点
     */
    public boolean contains(T key) {
        Node current = root;
        while (current != nil) {
            int cmp = key.compareTo(current.key);
            if (cmp == 0) {
                return true;  // 找到节点
            } else if (cmp < 0) {
                current = current.left;
            } else {
                current = current.right;
            }
        }
        return false;  // 未找到节点
    }

    /**
     *  Transplant操作：用v子树替换u子树
     */
    private void transplant(Node u, Node v) {
        if (u.parent == nil) {
            root = v;
        } else if (u == u.parent.left) {
            u.parent.left = v;
        } else {
            u.parent.right = v;
        }
        v.parent = u.parent;
    }

    /**
     * 找到以node为根的子树中的最小节点
     */
    private Node minimum(Node node) {
        while (node.left != nil) {
            node = node.left;
        }
        return node;
    }

    /**
     * 删除节点
     */
    public void delete(T key) {
        Node z = searchNode(key);
        if (z == nil) {
            return;  // 节点不存在
        }

        Node y = z;
        Node x;
        boolean yOriginalColor = y.color;

        // 情况1：z只有右子树或没有子树
        if (z.left == nil) {
            x = z.right;
            transplant(z, z.right);
        }
        // 情况2：z只有左子树
        else if (z.right == nil) {
            x = z.left;
            transplant(z, z.left);
        }
        // 情况3：z有两个子树
        else {
            y = minimum(z.right);  // 找到z的后继节点
            yOriginalColor = y.color;
            x = y.right;

            // 如果y是z的直接右子节点
            if (y.parent == z) {
                x.parent = y;
            } else {
                transplant(y, y.right);
                y.right = z.right;
                y.right.parent = y;
            }

            transplant(z, y);
            y.left = z.left;
            y.left.parent = y;
            y.color = z.color;
        }

        // 如果被删除的节点是黑色，需要修复红黑树性质
        if (yOriginalColor == BLACK) {
            deleteFixup(x);
        }
    }

    /**
     * 删除后修复红黑树性质
     */
    private void deleteFixup(Node x) {
        while (x != root && x.color == BLACK) {
            if (x == x.parent.left) {
                Node w = x.parent.right;  // 兄弟节点

                // 情况1：兄弟节点为红色
                if (w.color == RED) {
                    w.color = BLACK;
                    x.parent.color = RED;
                    leftRotate(x.parent);
                    w = x.parent.right;
                }

                // 情况2：兄弟节点的两个子节点都为黑色
                if (w.left.color == BLACK && w.right.color == BLACK) {
                    w.color = RED;
                    x = x.parent;
                } else {
                    // 情况3：兄弟节点的左子节点为红色，右子节点为黑色
                    if (w.right.color == BLACK) {
                        w.left.color = BLACK;
                        w.color = RED;
                        rightRotate(w);
                        w = x.parent.right;
                    }
                    // 情况4：兄弟节点的右子节点为红色
                    w.color = x.parent.color;
                    x.parent.color = BLACK;//交换颜色
                    w.right.color = BLACK;
                    leftRotate(x.parent);
                    x = root;  // 退出循环
                }
            } else {
                // 镜像情况：x是其父节点的右子节点
                Node w = x.parent.left;  // 兄弟节点

                // 情况1：兄弟节点为红色
                if (w.color == RED) {
                    w.color = BLACK;
                    x.parent.color = RED;
                    rightRotate(x.parent);
                    w = x.parent.left;
                }

                // 情况2：兄弟节点的两个子节点都为黑色
                if (w.right.color == BLACK && w.left.color == BLACK) {
                    w.color = RED;
                    x = x.parent;
                } else {
                    // 情况3：兄弟节点的右子节点为红色，左子节点为黑色
                    if (w.left.color == BLACK) {
                        w.right.color = BLACK;
                        w.color = RED;
                        leftRotate(w);
                        w = x.parent.left;
                    }
                    // 情况4：兄弟节点的左子节点为红色
                    w.color = x.parent.color;
                    x.parent.color = BLACK;
                    w.left.color = BLACK;
                    rightRotate(x.parent);
                    x = root;  // 退出循环
                }
            }
        }
        x.color = BLACK;  // 确保x为黑色
    }

    /**
     * 搜索节点（内部使用）
     */
    private Node searchNode(T key) {
        Node current = root;
        while (current != nil) {
            int cmp = key.compareTo(current.key);
            if (cmp == 0) {
                return current;  // 找到节点
            } else if (cmp < 0) {
                current = current.left;
            } else {
                current = current.right;
            }
        }
        return nil;  // 未找到节点
    }

    /**
     * 中序遍历，用于测试
     */
    public List<T> inorderTraversal() {
        List<T> result = new ArrayList<>();
        inorderHelper(root, result);
        return result;
    }

    private void inorderHelper(Node node, List<T> result) {
        if (node != nil) {
            inorderHelper(node.left, result);
            result.add(node.key);
            inorderHelper(node.right, result);
        }
    }

    /**
     * 打印树结构，用于可视化
     */
    public void printTree() {
        System.out.println("红黑树结构:");
        printHelper(root, "", true);
    }

    private void printHelper(Node node, String prefix, boolean isLeft) {
        if (node != nil) {
            System.out.println(prefix + (isLeft ? "├── " : "└── ") +
                    node.key + "(" + (node.color == RED ? "红" : "黑") + ")");
            printHelper(node.left, prefix + (isLeft ? "│   " : "    "), true);
            printHelper(node.right, prefix + (isLeft ? "│   " : "    "), false);
        }
    }

    // 测试方法
    public static void main(String[] args) {
        RedBlackTree<Integer> tree = new RedBlackTree<>();

        // 插入测试数据
//        int[] values = {10, 20, 30, 15, 5, 25, 35, 1, 6};
        //插入情况1:插入的新节点 N 是红黑树的根节点，这种情况下，我们把节点 N 的颜色由红色变为黑色，性质2（根是黑色）被满足。同时 N 被染成黑色后，红黑树所有路径上的黑色节点数量增加一个，性质5（从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点）仍然被满足。
//        Integer [] values = {1};
        //插入情况2：N的父节点是黑色，这种情况下，性质4（每个红色节点必须有两个黑色的子节点）和性质5没有受到影响，不需要调整。
//        Integer [] values ={1,2};
        //N 的父节点是红色（节点 P 为红色，其父节点必然为黑色），叔叔节点 U 也是红色。由于 P 和 N 均为红色，所有性质4被打破，此时需要进行调整。 这种情况下，先将 P 和 U 的颜色染成黑色，再将 G 的颜色染成红色。此时经过 G 的路径上的黑色节点数量不变，性质5仍然满足。但需要注意的是 G 被染成红色后，可能会和它的父节点形成连续的红色节点，此时需要递归向上调整。
//        Integer [] values ={10,5,15,3};//N=3
        Integer[] values = {16, 3, 7, 11, 9, 26, 18, 14, 15};
        for (Integer val : values) {
            tree.insert(val);
            System.out.println("插入 " + val + " 后:");
            tree.printTree();
            System.out.println("中序遍历: " + tree.inorderTraversal() + "\n");
        }

        // 查找测试
        Integer searchVal = 10;
        System.out.println("是否包含 " + searchVal + ": " + tree.contains(searchVal));

        // 删除测试
        Integer deleteVal = 20;
        tree.delete(deleteVal);
        System.out.println("删除 " + deleteVal + " 后:");
        tree.printTree();
        System.out.println("中序遍历: " + tree.inorderTraversal());
    }
}
